package entity;

/**
 * 红黑树
 *
 * @Author: wanqian
 * @Date: 2021/2/22 15:17
 */
public class RedBlackTree {

    private Node root;

    private int size;

    private static final boolean BLACK = false;
    private static final boolean RED = true;

    private class Node {
        public int id;

        private String value;

        private Node left;

        private Node right;

        private boolean color;

        public Node(int id, String value, Node left, Node right, boolean color) {
            this.id = id;
            this.value = value;
            this.left = left;
            this.right = right;
            this.color = color;
        }
    }

    public int getSize() {
        return size;
    }

    private boolean isRed(Node node) {
        if (node == null) {
            return BLACK;
        }

        return node.color;
    }

    /**
     * 左旋
     * 当节点右节点为红色需要左旋
     *
     * @param node
     * @return
     */
    private Node rotateLeft(Node node) {
        Node right = node.right;
        node.right = right.left;
        right.left = node;
        right.color = node.color;
        node.color = RED;
        return right;
    }

    /**
     * 右旋
     * 当左子节点为红色，左子节点的左子节点也为红色(3-4节点)
     *
     * @param node
     * @return
     */
    private Node rotateRight(Node node) {
        Node left = node.left;
        node.left = left.right;
        left.right = node;
        left.color = node.color;
        node.color = RED;
        return left;
    }

    /**
     * 颜色反转
     * 当左右节点都为红色时
     *
     * @param node
     */
    private void flipColors(Node node) {
        node.color = RED;
        node.left.color = BLACK;
        node.right.color = BLACK;
    }

    /**
     * 在整颗树上放入节点
     *
     * @param id
     * @param value
     */
    public void put(int id, String value) {
        root = put(root, id, value);
        //防止颜色反转导致根节点变色
        root.color = BLACK;
    }

    /**
     * 在节点上插入
     *
     * @param node
     * @param id
     * @param value
     * @return
     */
    private Node put(Node node, int id, String value) {
        if (node == null) {
            size++;
            return new Node(id, value, null, null, RED);
        }

        if (id < node.id) {
            node.left = put(node.left, id, value);
        } else if (id > node.id) {
            node.right = put(node.right, id, value);
        } else {
            node.value = value;
        }

        //右节点为红色，左节点为黑色
        if (isRed(node.right) && !isRed(node.left)) {
            node = rotateLeft(node);
        }

        //左节点为红色，左节点的左节点也为红色(3-4节点)
        if (isRed(node.left) && isRed(node.left.left)) {
            node = rotateRight(node);
        }

        //左右节点为红色
        if (isRed(node.left) && isRed(node.right)) {
            flipColors(node);
        }

        return node;
    }

    public String get(int id) {
        return get(root, id);
    }

    private String get(Node node, int id) {
        if (node == null) {
            return null;
        }

        if (id < node.id) {
            return get(node.left, id);
        } else if (id > node.id) {
            return get(node.right, id);
        } else {
            return node.value;
        }
    }

//    public static void main(String[] args) {
//        RedBlackTree redBlackTree = new RedBlackTree();
//        redBlackTree.put(1, "张三");
//        redBlackTree.put(2, "李四");
//        redBlackTree.put(3, "王五");
//
//        System.out.println(redBlackTree.getSize());
//        System.out.println(redBlackTree.get(3));
//        System.out.println(redBlackTree.get(2));
//        System.out.println(redBlackTree.get(1));
//        System.out.println(redBlackTree.get(4));
//
//    }
}
